package edu.nepu.bigdata.kafka;

import edu.nepu.bigdata.compile.DynamicCompileAndRun;
import edu.nepu.bigdata.utils.CommandLineParser;
import edu.nepu.bigdata.utils.TimeUtils;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.text.ParseException;
import java.util.*;

public class ConsumerKafka {
    public static void main(String[] args) throws Exception {
        Logger logger = LoggerFactory.getLogger(ConsumerKafka.class);
        CommandLineParser.parseKafkaArgs(args);
        String topic = CommandLineParser.getArg("topic");

        Properties properties = new Properties();
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, CommandLineParser.getArg("bootstrap.servers"));
        properties.put(ConsumerConfig.GROUP_ID_CONFIG, "my_tool");
        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(properties);
        // 订阅主题
        //consumer.subscribe(Collections.singletonList(topic));

        // 首先获取分区的当前结束偏移量。如果我们无法根据时间戳找到合适的偏移量，即满足时间戳要求的消息尚未发送到 Kafka，则将使用此方法，在这种情况下，我们只使用最新的偏移量。
        // 我们需要在按时间查询偏移量之前获取最新的偏移量，以确保不会错过任何消息。
        if (CommandLineParser.getArg("timestamp") != null || CommandLineParser.getArg("date_hms") != null) {
            final long timestamp;
            if (CommandLineParser.getArg("date_hms") != null){
                timestamp = TimeUtils.stringToTimestamp(CommandLineParser.getArg("date_hms"), "yyyy-MM-dd-HH:mm:ss");
            } else {
                timestamp = Long.parseLong(CommandLineParser.getArg("timestamp"));
            }
            //consumer.poll(1000);
            List<PartitionInfo> partitionInfos = consumer.partitionsFor(topic);
            ArrayList<TopicPartition> topicPartitions = new ArrayList<>();
            for (PartitionInfo partitionInfo : partitionInfos) {
                topicPartitions.add(new TopicPartition(partitionInfo.topic(), partitionInfo.partition()));
            }
            consumer.assign(topicPartitions);
            Map<TopicPartition, Long> topicPartitionLongMap = consumer.endOffsets(topicPartitions);
            Map<TopicPartition, Long> startingTimestamps = new HashMap<>();
            Map<TopicPartition, Long> initialOffsets = new HashMap<>();
            topicPartitions.forEach(tp -> startingTimestamps.put(tp, timestamp));
            Map<TopicPartition, OffsetAndTimestamp> topicPartitionOffsetAndTimestampMap = consumer.offsetsForTimes(startingTimestamps);
            for (TopicPartition tp : topicPartitions) {
                // offset may not have been resolved
                if (topicPartitionOffsetAndTimestampMap.containsKey(tp) && topicPartitionOffsetAndTimestampMap.get(tp) != null) {
                    initialOffsets.put(tp, topicPartitionOffsetAndTimestampMap.get(tp).offset());
                } else {
                    initialOffsets.put(tp, topicPartitionLongMap.get(tp));
                }
            }
            initialOffsets.forEach(consumer::seek);
        } else {
            logger.error("请传入指定消费的时间戳");
        }
        try {
            DynamicCompileAndRun dynamicCompileAndRun = new DynamicCompileAndRun(CommandLineParser.getArg("fullClassName"), CommandLineParser.getArg("methodName"));
            while (true) {
                consumer.poll(100).forEach(record -> {
                    try {
                        dynamicCompileAndRun.run(record.partition(), record.offset(), record.key(), record.value());
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                });
            }
        } finally {
            consumer.close();
        }
    }
}
